﻿#region LicenseHeader

// Copyright 2012 The Trustees of Indiana University.  All rights reserved.
// 
// The Indiana University licenses this file to you under the Apache License, 
// Version 2.0 (the "License"); you may not use this file except in compliance 
// with the License.  You may obtain a copy of the License at 
// 
// http://www.apache.org/licenses/LICENSE-2.0
// 
// Unless required by applicable law or agreed to in writing, software 
// distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
//  WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.  
// See the License for the specific language governing permissions and 
// limitations under the License.
// 
// @author Thilina Gunarathne (tgunarat@indiana.edu)

#endregion

using System;
using System.Collections.Generic;
using System.Linq;
using System.Threading;
using AzureMRCore.DataModel;
using Microsoft.WindowsAzure;
using Microsoft.WindowsAzure.StorageClient;

namespace AzureMRCore.MapRed
{
    /// <summary>
    ///   Helper class for ReduceWorker. Acts as the configuration holder for Reduce tasks. Handles all the Table related functionality for Reduce Tasks.
    /// </summary>
    public class ReduceContext
    {
        private bool _isAllMapDone;

        public ReduceContext(Dictionary<string, object> reduceParams, CloudStorageAccount storageAccount,
                             CloudTableClient cloudTableClient, string tableBaseAddress)
        {
            JobID = (string) reduceParams[Constants.JOBID];
            AppName = (string) reduceParams[Constants.MR_APP_NAME];
            ProgramParams = (string) reduceParams[Constants.PROGRAM_PARAMS];
            ReduceID = (string) reduceParams[Constants.REDUCEID];
            OutDir = ((string) reduceParams[Constants.OUT_DIR]).ToLower();

            NumMaps = Int32.Parse((string) reduceParams[Constants.NUM_MAPS]);
            Iteration = (string) reduceParams[Constants.ITERATION];
            IsTableBCast = false;

            ReduceInputDataTableName = (string) reduceParams[Constants.TASK_TABLE_NAME];
            cloudTableClient.CreateTableIfNotExist(ReduceInputDataTableName);
            ReduceInputDataContext = new ReduceInputDataContext(tableBaseAddress,
                                                                storageAccount.Credentials, ReduceInputDataTableName);

            MapContext = new MapMetaDataContext(tableBaseAddress, storageAccount.Credentials,
                                                DataModelUtils.GetMapTableName(JobID));

            string[] strings = ReduceID.Split('_');
            ReducerIndex = Int32.Parse(strings[0]);

            int retry = 0;
            RetryLabel:
            try
            {
                //wait for the things to start up as tables might not be loaded yet.
                while (WaitForMapTaskStart())
                {
                    Thread.Sleep(100);
                }
                CreateReduceTables(storageAccount, cloudTableClient, tableBaseAddress);
            }
            catch (Exception e)
            {
                if (retry < Constants.NUM_RETRIES)
                {
                    retry++;
                    Thread.Sleep(100);
                    goto RetryLabel;
                }
                throw new Exception("Reduce Context : TablecreateLabel failed. Retried " + retry, e);
            }
        }

        public string ProgramParams { get; set; }
        public string AppName { get; set; }
        public string JobID { get; set; }
        public string ReduceID { get; set; }
        public string OutDir { get; set; }
        public MapMetaDataContext MapContext { get; set; }
        public int NumMaps { get; set; }
        public string ReduceInputDataTableName { get; set; }
        public int Attempt { get; set; }
        public string Iteration { get; set; }
        public int ReducerIndex { get; set; }
        public bool IsTableBCast { get; set; }

        public ReduceInputDataContext ReduceInputDataContext { get; set; }
        public ReduceDataModel ReduceTask { get; set; }
        public ReduceMetaDataContext ReduceDataContext { get; set; }

        private void CreateReduceTables(CloudStorageAccount storageAccount, CloudTableClient cloudTableClient,
                                        string tableBaseAddress)
        {
            ReduceDataContext = new ReduceMetaDataContext(tableBaseAddress,
                                                          storageAccount.Credentials,
                                                          DataModelUtils.GetReduceTableName(JobID));
            int retryCount = 0;
            while (!cloudTableClient.DoesTableExist(ReduceDataContext.TableName) & retryCount < 20)
            {
                Thread.Sleep(100);
                retryCount++;
            }

            IQueryable<ReduceDataModel> reduceTaskEnum =
                ReduceDataContext.ReduceTasks.Where(p => (p.RowKey == ReduceID) & (p.PartitionKey == Iteration));
            IEnumerator<ReduceDataModel> reduceTaskEnumerator = reduceTaskEnum.GetEnumerator();
            if (!reduceTaskEnumerator.MoveNext())
            {
                throw new Exception("No table entry for the reduce task");
            }
            ReduceTask = reduceTaskEnumerator.Current;
        }

        // TODO : try to get rid of this
        private bool WaitForMapTaskStart()
        {
            IQueryable mapTasks = MapContext.MapTasks.Take(1);
            return !(mapTasks.GetEnumerator().MoveNext());
        }


        public int SetStateToExecuting(string roleID)
        {
            int attempt;
            ReduceTask.TaskStatus = MapDataModel.EXECUTING;
            if (ReduceTask.Attempt1Status == MapDataModel.SUBMITTED)
            {
                ReduceTask.Attempt1Status = MapDataModel.EXECUTING;
                ReduceTask.Attempt1StartTime = string.Format("{0:G}", DateTime.Now);
                attempt = 1;
            }
            else if (ReduceTask.Attempt2Status == MapDataModel.SUBMITTED)
            {
                ReduceTask.Attempt2Status = MapDataModel.EXECUTING;
                ReduceTask.Attempt2StartTime = string.Format("{0:G}", DateTime.Now);
                attempt = 2;
            }
            else if (ReduceTask.Attempt3Status == MapDataModel.SUBMITTED)
            {
                ReduceTask.Attempt3Status = MapDataModel.EXECUTING;
                ReduceTask.Attempt3StartTime = string.Format("{0:G}", DateTime.Now);
                attempt = 3;
            }
            else
            {
                //to do handle job failure
                attempt = -1;
                ReduceTask.TaskStatus = MapDataModel.FAILED;
            }

            ReduceTask.WorkerRole = roleID;
            ReduceDataContext.UpdateObject(ReduceTask);
            ReduceDataContext.SaveChangesWithRetries();
            return attempt;
        }

        internal void UpdateReduceTaskStatus(int reduceTaskStatus)
        {
            switch (Attempt)
            {
                case 1:
                    ReduceTask.Attempt1Status = reduceTaskStatus;
                    break;
                case 2:
                    ReduceTask.Attempt2Status = reduceTaskStatus;
                    break;
                case 3:
                    ReduceTask.Attempt3Status = reduceTaskStatus;
                    break;
            }
            if (reduceTaskStatus == MapDataModel.FINISHED)
            {
                ReduceTask.EndTime = string.Format("{0:G}", DateTime.Now);
                ReduceTask.TaskStatus = MapDataModel.FINISHED;
            }
            ReduceDataContext.UpdateObject(ReduceTask);
            ReduceDataContext.SaveChangesWithRetries();
        }

        internal bool IsDoneFetching(int numMapTasks, string iteration, string appName, int pollingInterval,
                                     int fetchedBLOBCount)
        {
            if (!_isAllMapDone)
            {
                if (!IsAllMapDone(numMapTasks, iteration, appName, pollingInterval)) return false;
                _isAllMapDone = true;
            }

            IQueryable<MapDataModel> mapTasks =
                MapContext.MapTasks.Where(p => p.PartitionKey == iteration & p.AppName == AppName).AsTableServiceQuery();
            int totalDataProducts = Enumerable.Aggregate(mapTasks, 0,
                                                         (current, mapTask) =>
                                                         current + mapTask.ReduceProductCounts[ReducerIndex]);

//            Trace.WriteLine("Reduce IsDoneFectching. MapOutCount:" + totalDataProducts + " FetchedRecordsCount:" +
//                            fetchedBLOBCount + " Reduce ID:" + ReduceID + " i:" + Iteration);
            return fetchedBLOBCount == totalDataProducts;
        }

        public bool IsAllMapDone(int numMapTasks, string iteration, string appName, int pollingInterval)
        {
            IQueryable<MapDataModel> executingMapTasks =
                MapContext.MapTasks.Where(
                    p => p.TaskStatus != MapDataModel.FINISHED & p.PartitionKey == iteration & p.AppName == appName)
                    .Take(1);

            //Done processing all the map tasks)
            if (Enumerable.Count(executingMapTasks) > 0)
            {
                Thread.Sleep(pollingInterval/5);
                return false;
            }

            //hnadling the case where the MapTasks are still not populated in to the table
            CloudTableQuery<MapDataModel> allMapTasks =
                MapContext.MapTasks.Where(p => p.PartitionKey == iteration & p.AppName == appName).
                    AsTableServiceQuery();
            if (Enumerable.Count(allMapTasks) != numMapTasks)
            {
                Thread.Sleep(pollingInterval/10);
                return false;
            }

            IQueryable<MapDataModel> doubleCheckMapTasks =
                allMapTasks.Where(p => p.TaskStatus != MapDataModel.FINISHED).Take(1);

            //Done processing all the map tasks)
            if (Enumerable.Count(doubleCheckMapTasks) > 0)
            {
                Thread.Sleep(pollingInterval/5);
                return false;
            }
            return true;
        }
    }
}